import onnxruntime as rt
import BaseData
import numpy as np
import cv2

tag = ['0faya', 'faya']
sess = rt.InferenceSession('out_file/tudou.onnx', None)
input_name = sess.get_inputs()[0].name
out_name = sess.get_outputs()[0].name

cap = cv2.VideoCapture(0)
ret_flag,Vshow = cap.read()
dt = BaseData.ImageData(Vshow, backbone="MobileNet")
input_data = dt.to_tensor()

pred_onx = sess.run([out_name], {input_name: input_data})
ort_output = pred_onx[0]
idx = np.argmax(ort_output, axis=1)[0]
print('result:' + tag[idx])
